fix: exclude FP8 quantized weights from MXFP4 de-osc by hann-wang · Pull Request #1081 · AMD-AGI/Primus

hann-wang · 2026-09-03T01:13:50Z

What's wrong

MXFP4 weight-deosc process selects eligible weights by looking for quantized_weight_buffer. FP4 and FP8 quantized modules both have this buffer when gradient accumulation enabled.

What's changed

MXFP4 quantized module creates an empty buffer in float4_e2m1fn_x2 dtype when gradient accumulation disabled.
Check the dtype of quantized_weight_buffer in MXFP4 weight de-osc.

@HuangWei-95

Copilot

🔵 Needs a closer look

It affects low-level quantization/training behavior via a custom FP4 dtype marker and should get final human validation in a real runtime/config matrix before approval.

Pull request overview

This PR refines MXFP4 weight de-oscillation eligibility detection so FP8-quantized modules are no longer mistakenly picked up when quantized_weight_buffer exists for both FP4 and FP8 paths (e.g., when gradient accumulation is enabled).

Changes:

In Primus-Turbo FP4 paths with get_num_microbatches() == 1, set quantized_weight_buffer to an empty tensor with float4_e2m1fn_x2 dtype (instead of torch.uint8) to act as an FP4-specific runtime marker.
In the weight de-osc eligibility scan, require quantized_weight_buffer.dtype == float4_e2m1fn_x2 to filter out FP8 buffers.

File summaries

File	Description
`primus/backends/megatron/core/optimizer/weight_deosc.py`	Tightens FP4 eligibility detection by checking `quantized_weight_buffer` dtype.
`primus/backends/megatron/core/extensions/primus_turbo.py`	Makes the FP4 “marker” buffer use `float4_e2m1fn_x2` dtype so de-osc can reliably identify FP4 modules.

Review details

Files reviewed: 2/2 changed files
Comments generated: 0
Review effort level: Lite

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

Copilot

🟡 Changes recommended

There are concrete runtime-breaking issues (undefined weights variable in multiple FP4 paths) and the new dtype check can raise/disable eligibility when the MXFP4 dtype symbol is unavailable or the marker isn’t a tensor.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Review details

Suppressed comments (3)

primus/backends/megatron/core/extensions/primus_turbo.py:1365

weights is not defined in this forward_internal (the local variable is weight), so device=weights.device will raise NameError the first time this FP4 path runs.

                        self.quantized_weight_buffer = torch.empty(
                            0, device=weights.device, dtype=float4_e2m1fn_x2
                        )

primus/backends/megatron/core/extensions/primus_turbo.py:1571

weights is not defined in this forward_internal (the local variable is weight), so device=weights.device will raise NameError the first time this FP4 path runs.

                        self.quantized_weight_buffer = torch.empty(
                            0, device=weights.device, dtype=float4_e2m1fn_x2
                        )

primus/backends/megatron/core/extensions/primus_turbo.py:1790

weights is not defined in this forward_internal (the local variable is weight), so device=weights.device will raise NameError the first time this FP4 path runs.

                        self.quantized_weight_buffer = torch.empty(
                            0, device=weights.device, dtype=float4_e2m1fn_x2
                        )

Files reviewed: 2/2 changed files
Comments generated: 2
Review effort level: Lite

+                quantized_weight_buffer = getattr(module, "quantized_weight_buffer", None)
+                if quantized_weight_buffer is None or quantized_weight_buffer.dtype != _float4_e2m1fn_x2:
                    continue


Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>

Copilot

🔵 Needs a closer look

Several FP4 single-microbatch code paths reference weights.device where weights is undefined, which will raise a runtime NameError.

Review details

Suppressed comments (3)

primus/backends/megatron/core/extensions/primus_turbo.py:1571

weights is not defined in this scope; using weights.device will raise a NameError on the FP4 single-microbatch path. Use weight.device here.

                        self.quantized_weight_buffer = torch.empty(
                            0, device=weights.device, dtype=float4_e2m1fn_x2
                        )

primus/backends/megatron/core/extensions/primus_turbo.py:1790

weights is not defined in this scope; using weights.device will raise a NameError on the FP4 single-microbatch path. Use weight.device here.

                        self.quantized_weight_buffer = torch.empty(
                            0, device=weights.device, dtype=float4_e2m1fn_x2
                        )

primus/backends/megatron/core/extensions/primus_turbo.py:1365

weights is not defined in this scope; using weights.device will raise a NameError the first time this FP4 single-microbatch path runs. Use the local weight parameter's device instead.

                        self.quantized_weight_buffer = torch.empty(
                            0, device=weights.device, dtype=float4_e2m1fn_x2
                        )

Files reviewed: 2/2 changed files
Comments generated: 0 new
Review effort level: Lite

Copilot

🔵 Needs a closer look

The new dtype-based eligibility check can raise AttributeError when quantized_weight_buffer is not a tensor-like object exposing .dtype, so it should be made defensive before merging.

Review details

Suppressed comments (1)

primus/backends/megatron/core/optimizer/weight_deosc.py:355

_build_eligible_ids now unconditionally accesses quantized_weight_buffer.dtype. Some callers (including existing unit tests) use non-tensor markers for quantized_weight_buffer, which would raise AttributeError here. Consider reading dtype defensively (and update tests/markers to use the FP4 dtype marker).

        match weights the forward actually quantizes.
        """

Files reviewed: 2/2 changed files
Comments generated: 0 new
Review effort level: Lite

Copilot

🔵 Needs a closer look

The new dtype-based eligibility check can raise an AttributeError for non-tensor markers (and will break existing CPU unit tests) unless guarded.

Review details

Suppressed comments (1)

primus/backends/megatron/core/optimizer/weight_deosc.py:366

_build_eligible_ids now assumes quantized_weight_buffer always has a .dtype attribute; if a module (or the CPU unit tests in tests/unit_tests/backends/megatron/test_weight_deosc.py) uses a sentinel object to mark “FP4 forward ran”, this will raise AttributeError at runtime. Use getattr(..., "dtype", None) to keep the dtype filter without crashing on non-tensor sentinels.

                quantized_weight_buffer = getattr(module, "quantized_weight_buffer", None)
                if quantized_weight_buffer is None or quantized_weight_buffer.dtype != _float4_e2m1fn_x2:

Files reviewed: 2/2 changed files
Comments generated: 0 new
Review effort level: Lite

fix: exclude FP8 quantized weights from MXFP4 de-osc

4be54cb

Copilot AI lite review requested due to automatic review settings September 3, 2026 01:13

hann-wang requested review from Xiaoming-AMD, limou102 and wenxie-amd as code owners September 3, 2026 01:13

Copilot started reviewing on behalf of hann-wang September 3, 2026 01:14 View session

Copilot AI reviewed Sep 3, 2026

View reviewed changes

chore: lint

e5286cb

Copilot AI review requested due to automatic review settings September 3, 2026 01:23

Copilot started reviewing on behalf of hann-wang September 3, 2026 01:24 View session

Copilot AI reviewed Sep 3, 2026

View reviewed changes

Fix typo in variable name for weight buffer

47f24a8

Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>

Copilot AI review requested due to automatic review settings September 9, 2026 02:18

Copilot started reviewing on behalf of hann-wang September 9, 2026 02:19 View session

Copilot AI reviewed Sep 9, 2026

View reviewed changes

fix: typo

11c9ebe

Copilot AI review requested due to automatic review settings September 9, 2026 02:23

Copilot started reviewing on behalf of hann-wang September 9, 2026 02:23 View session

Merge branch 'main' into dev/han/fix-mxfp4-deosc-with-fp8-weights

6ffd82e

Copilot AI reviewed Sep 9, 2026

View reviewed changes

Copilot AI review requested due to automatic review settings September 9, 2026 02:27

Copilot started reviewing on behalf of hann-wang September 9, 2026 02:28 View session

Copilot AI reviewed Sep 9, 2026

View reviewed changes

Navigation Menu

Uh oh!

fix: exclude FP8 quantized weights from MXFP4 de-osc - #1081

fix: exclude FP8 quantized weights from MXFP4 de-osc#1081
hann-wang wants to merge 5 commits into
AMD-AGI:mainfrom
hann-wang:dev/han/fix-mxfp4-deosc-with-fp8-weights

hann-wang commented Sep 3, 2026

Uh oh!

Copilot AI left a comment

Uh oh!

Copilot AI left a comment

Uh oh!

Uh oh!

Copilot AI left a comment

Uh oh!

Copilot AI left a comment

Uh oh!

Copilot AI left a comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Uh oh!

Conversation

hann-wang commented Sep 3, 2026

What's wrong

What's changed

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

🔵 Needs a closer look

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

🟡 Changes recommended

Suppressed comments (3)

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

🔵 Needs a closer look

Suppressed comments (3)

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

🔵 Needs a closer look

Suppressed comments (1)

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

🔵 Needs a closer look

Suppressed comments (1)

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants